{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Tutorial 6: Train NicheTrans* on STARmap PLUS data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os, time, datetime, warnings\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.optim import lr_scheduler\n",
    "\n",
    "from model.nicheTrans_ct import *\n",
    "from datasets.data_manager_STARmap_PLUS import AD_Mouse\n",
    "\n",
    "from utils.utils import *\n",
    "from utils.utils_training_STARmap_PLUS import train, test\n",
    "from utils.utils_dataloader import *\n",
    "\n",
    "warnings.filterwarnings(\"ignore\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Initialize the args and fix seeds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "==========\n",
      "Args:Namespace(AD_adata_path='/home/wzk/ST_data/AD_mouse2/norm/AD_mouses_adata', Wild_type_adata_path='/home/wzk/ST_data/AD_mouse2/norm/benigh_mouses', dropout_rate=0.25, eval_step=1, gamma=0.1, gpu_devices='0', label_path='/home/wzk/ST_data/AD_mouse/generated_by_zhikang_2_filtering_tao15_abeta50', lr=0.0001, max_epoch=20, n_top_genes=2000, noise_rate=0.5, optimizer='adam', save_dir='./log', seed=1, stepsize=20, test_batch=32, train_batch=128, weight_decay=0.0005, workers=4)\n",
      "==========\n"
     ]
    }
   ],
   "source": [
    "%run ./args/args_STARmap_PLUS.py\n",
    "args = args\n",
    "\n",
    "set_seed(args.seed)\n",
    "os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_devices\n",
    "\n",
    "print(\"==========\\nArgs:{}\\n==========\".format(args))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Initialize dataloaders and NicheTrans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "------Calculating spatial graph...\n",
      "The graph contains 124464 edges, 10372 cells.\n",
      "12.0000 neighbors per cell on average.\n",
      "------Calculating spatial graph...\n",
      "The graph contains 115608 edges, 9634 cells.\n",
      "12.0000 neighbors per cell on average.\n",
      "------Calculating spatial graph...\n",
      "The graph contains 96408 edges, 8034 cells.\n",
      "12.0000 neighbors per cell on average.\n",
      "=> AD Mouse loaded\n",
      "Dataset statistics:\n",
      "  ------------------------------\n",
      "  subset   | # num | \n",
      "  ------------------------------\n",
      "  train    |  10372 spots, 894.0 positive tao, 291.0 positive plaque \n",
      "  test     |   9634 spots, 620.0 positive tao, 195.0 positive plaque \n",
      "  ------------------------------\n"
     ]
    }
   ],
   "source": [
    "# create the dataloaders\n",
    "dataset = AD_Mouse(AD_adata_path=args.AD_adata_path, Wild_type_adata_path=args.Wild_type_adata_path, label_path=args.label_path, n_top_genes=args.n_top_genes)\n",
    "trainloader, testloader = ad_mouse_dataloader(args, dataset)\n",
    "\n",
    "# create the model\n",
    "source_dimension, target_dimension = dataset.rna_length, dataset.target_length\n",
    "model = NicheTrans_ct(source_length=source_dimension, target_length=target_dimension, noise_rate=args.noise_rate, dropout_rate=args.dropout_rate)\n",
    "model = nn.DataParallel(model).cuda()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Initialize loss function (criterion) and optimizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "criterion = nn.MSELoss()\n",
    "\n",
    "if args.optimizer == 'adam':\n",
    "    optimizer = torch.optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay)\n",
    "elif args.optimizer == 'SGD':\n",
    "    optimizer = torch.optim.SGD(model.parameters(), lr=args.lr)\n",
    "else:\n",
    "    print('unexpected optimizer')\n",
    "\n",
    "if args.stepsize > 0:\n",
    "    scheduler = lr_scheduler.StepLR(optimizer, step_size=args.stepsize, gamma=args.gamma)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model training and testing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "start_time = time.time()\n",
    "\n",
    "for epoch in range(args.max_epoch):\n",
    "    last_epoch = epoch + 1 == args.max_epoch\n",
    "\n",
    "    print(\"==> Epoch {}/{}\".format(epoch+1, args.max_epoch))\n",
    "    \n",
    "    ################\n",
    "    train(args, model, criterion, optimizer, trainloader, dataset.target_panel, ct_information=True)\n",
    "    if args.stepsize > 0: scheduler.step()\n",
    "    \n",
    "    if (epoch+1) % args.eval_step == 0:\n",
    "        pearson = test(args, model, testloader, dataset.target_panel, last_epoch, ct_information=True)\n",
    "\n",
    "    if last_epoch==True:\n",
    "        torch.save(model.state_dict(), 'NicheTrans_*_STARmap_PLUS.pth')\n",
    "    ################\n",
    "\n",
    "elapsed = round(time.time() - start_time)\n",
    "elapsed = str(datetime.timedelta(seconds=elapsed))\n",
    "print(\"Finished. Total elapsed time (h:m:s): {}\".format(elapsed))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pytorch_zk",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.19"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}